/*
     Created by Paul Marinescu and George Candea
     Copyright (C) 2009 EPFL (Ecole Polytechnique Federale de Lausanne)

     This file is part of LFI (Library-level Fault Injector).

     LFI is free software: you can redistribute it and/or modify it  
     under the terms of the GNU General Public License as published by the  
     Free Software Foundation, either version 3 of the License, or (at  
     your option) any later version.

     LFI is distributed in the hope that it will be useful, but  
     WITHOUT ANY WARRANTY; without even the implied warranty of  
     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU  
     General Public License for more details.

     You should have received a copy of the GNU General Public  
     License along with LFI. If not, see http://www.gnu.org/licenses/.

     EPFL
     Dependable Systems Lab (DSLAB)
     Room 330, Station 14
     1015 Lausanne
     Switzerland
*/

#include <errno.h>

#define _GNU_SOURCE
#define __USE_GNU

#include <dlfcn.h>
#include <execinfo.h>
#include <time.h>
/* the maximum number of frames in a stack trace */
#define TRACE_SIZE	100
/* human readable log file (overwritten at each run) */
#define LOGFILE		"inject.log"
#define LOGGING		0
/* machine readable log file used for injection replay (overwritten at each run) */
#define	REPLAYFILE	"replay.xml"
#define MAPFILE		"/proc/self/maps"
#define MAXINJECT	2000000

/* only used to enchance readbility */
#ifndef __in
#define __in
#endif

#ifndef __out
#define __out
#endif

#ifndef NULL
#define NULL	(void*)0
#endif

/* structures used to keep fault information in random injection mode */
struct error_description
{
	int return_value;
	int errno_value;
};

struct fninfov1
{
	char function_name[256];
	int err_count;
	struct error_description* errors;
};
/* structure used to keep fault information in plan-injection mode */
struct fninfov2
{
	char function_name[256];
	int call_count;
	int return_value;
	int errno_value;
	int call_original;
};

static __thread long return_address;
static __thread int return_code;
static __thread int return_errno;
static __thread int no_intercept;

/*
   avoid including the standard headers because the compiler will likely
   generate warnings when injecting faults in an already declared function.
   However, this won't allow us to inject faults in these functions
   ALTERNATIVE: use dlsym to get pointers to the functions
*/
extern void srand(unsigned int _Seed);
extern int rand(void);
int printf(const char * _Format, ...);

void determine_actionv1(__in char* fn, int call_count,
						__out int* call_original, __out int *return_error, __out int* return_code, __out int* return_errno);

void determine_actionv1_o(struct fninfov1* fn_details, __in char* function_name, int call_count, __out int* call_original, __out int *return_error, __out int* return_code, __out int* return_errno);


void determine_actionv2(__in char* fn, __in int call_count,
						__out int* call_original, __out int *return_error, __out int* return_code, __out int* return_errno);
int find_range();

/************************************************************************/
/*	GENERATE_STUBv2 - macro to generate stub functions targeted for x86 */
/*	UNSAFE: clobbered non-volatile registeres may not be restored       */
/*	  when `forcing` an exit (i.e. when doing a leave/ret or leave/jmp) */
/*	SOLUTION: manually push/pop all non-volatile registers OR           */
/*	  check assembly listing generated by the compiler to determine if  */
/*	  any registers need to be `pop`-ed (per-compiler solution)         */
/************************************************************************/
#define GENERATE_STUBv2(FUNCTION_NAME) \
	void FUNCTION_NAME (void) \
{ \
	int nptrs; \
	void* buffer[TRACE_SIZE]; \
	char** sym; \
	int call_original, return_error; \
	static void * (*original_fn_ptr)(); \
	/* we can't call write directly because it would prevent us for injecting faults in `write`
       (injecting requires the creation of a function with the same name but the prototype is
	   different). We, use dlsym instead
	*/ \
	static int * (*original_write_ptr)(int, void*, int); \
	static int call_count = 1; \
	int initial_no_intercept; \
	\
	/* defaults */ \
	call_original = 1; \
	return_error = 0; \
	return_code = 0; \
	return_errno = 0; \
	nptrs = 0; \
	\
	initial_no_intercept = no_intercept; \
	if (no_intercept || 0 == init_done  /* || injx_count >= MAXINJECT  || time(NULL) - tt < 2 */) { \
		call_original = 1; \
		return_error = 0; \
	} \
	else {\
		no_intercept = 1; \
		nptrs = backtrace(buffer, TRACE_SIZE); \
		if (0 == exe_start) { \
			find_range(); \
		} \
		if (buffer[1] < exe_end && buffer[1] > exe_start ) /* called by our target */ \
		{ \
			/* sym = backtrace_symbols(&buffer[1], 1); \
			if (sym) {\				
				free(sym); \
			} */ \
			\
			determine_action(&function_info_ ## FUNCTION_NAME, #FUNCTION_NAME, call_count, &call_original, &return_error, &return_code, &return_errno); \
			\
			/* if (0 == return_code) { return_error = 0; call_original = 1; } */\
			++call_count; \
			if (return_error || (!call_original)) \
			{ \
				if(!original_write_ptr) \
				original_write_ptr = (int *(*)(int, void*, int)) dlsym(RTLD_NEXT, "write"); \
				original_write_ptr(log_fd, "stacktrace follows:\n", 20); \
				backtrace_symbols_fd(buffer, nptrs, log_fd); \
				/* printf("injecting at %s\n", #FUNCTION_NAME); */\
				++injx_count; \
			} \
		} else { \
			call_original = 1; \
			return_error = 0; \
			/* printf("Not called by target %x %x - %x\n", buffer[1], exe_start, exe_end); */ \
		}\
	} \
	/* printf("done intercepted %s\n", #FUNCTION_NAME); */ \
	\
	if(!original_fn_ptr) \
		original_fn_ptr = (void *(*)()) dlsym(RTLD_NEXT, #FUNCTION_NAME); \
	\
	if(!original_fn_ptr) \
		printf("unable to get original ptr\n"); \
	\
	no_intercept = initial_no_intercept; \
	/* disabled - unlikely to be useful in practice */ \
	if (0 && call_original && return_error) \
	{ \
	/* save the original return value */ \
	__asm__ ("movl 0x4(%%ebp), %%eax;" : "=a"(return_address)); \
	\
	__asm__ ("leave"); \
	__asm__ ("addl $0x4, %esp"); \
	/* at this point the stack is gone */ \
	\
	/* make the call the original function with the same stack */ \
	__asm__ ("call *%%eax;" : : "a"(original_fn_ptr)); \
	\
	errno = return_errno; \
	\
	/* push back the original return value */ \
	__asm__ ("pushl %%eax;" : : "a"(return_address)); \
	\
	__asm__ ("ret" : : "a"(return_code)); \
} \
	else if (return_error) \
{ \
	errno = return_errno; \
	__asm__ ("nop" : : "a"(return_code)); \
	return; \
} \
	else if (call_original) \
{ \
	__asm__ ("nop" : : "a"(original_fn_ptr)); \
	__asm__ ("add $0x1d4, %esp"); /* magic - compiler dependent */ \
	__asm__ ("pop %ebx"); \
	\
	__asm__ ("pop %ebp"); \
	__asm__ ("jmp *%eax"); \
} \
}

/************************************************************************/
/*	GENERATE_STUBv2_x64 - macro to generate stub functions on x64       */
/*	UNSAFE: can present undefined behaviour when an exception is        */
/*         encountered here or any called function                      */
/*	(but it's ok if an exception happens in the original function)      */
/*	PROBLEM: the x64 ABI has a clear notion of function prologue and    */
/*                                                   epilogue but...    */
/*		we are using the stack to save/restore the initial state        */
/*      (the one where we were called), i.e. register values,           */
/*      outside the prologue/epilogue                                   */
/*		http://msdn.microsoft.com/en-us/library/8ydc79k6(VS.80).aspx    */
/************************************************************************/

#define GENERATE_STUBv2_x64(FUNCTION_NAME) \
	void FUNCTION_NAME (void) \
{ \
	int nptrs; \
	void* buffer[TRACE_SIZE]; \
	int call_original, return_error; \
	static void * (*original_fn_ptr)(); \
	/* we can't call write directly because it would prevent us for injecting faults in `write`
       (injecting requires the creation of a function with the same name but the prototype is
	   different). We, use dlsym instead
	*/ \
	static int * (*original_write_ptr)(int, void*, int); \
	static int call_count = 1; \
	int initial_no_intercept; \
	\
	__asm__ ("push %r13"); \
	__asm__ ("push %r12"); \
	__asm__ ("push %r11"); \
	__asm__ ("push %r10"); \
	__asm__ ("push %rdi"); \
	__asm__ ("push %rsi"); \
	__asm__ ("push %rbx"); \
	__asm__ ("push %rcx"); \
	__asm__ ("push %rdx"); \
	__asm__ ("push %r8"); \
	__asm__ ("push %r9"); \
	/* defaults */ \
	call_original = 1; \
	return_error = 0; \
	return_code = 0; \
	return_errno = 0; \
	nptrs = 0; \
	/* printf("intercepted %s\n", #FUNCTION_NAME); */\
	\
	initial_no_intercept = no_intercept; \
	if (no_intercept || 0 == init_done) { \
		call_original = 1; \
		return_error = 0; \
	} \
	else {\
		no_intercept = 1; \
		nptrs = backtrace(buffer, TRACE_SIZE); \
		/* determine_action(#FUNCTION_NAME, call_count, &call_original, &return_error, &return_code, &return_errno); */ \
		determine_action_o(function_infov1_ ## FUNCTION_NAME, call_count, &call_original, &return_error, &return_code, &return_errno); \
		++call_count; \
		if (return_error || (!call_original)) \
		{ \
			if(!original_write_ptr) \
				original_write_ptr = (int *(*)(int, void*, int)) dlsym(RTLD_NEXT, "write"); \
			original_write_ptr(log_fd, "stacktrace follows:\n", 20); \
			backtrace_symbols_fd(buffer, nptrs, log_fd); \
		} \
	} \
	\
	if(!original_fn_ptr) \
		original_fn_ptr = (void *(*)()) dlsym(RTLD_NEXT, #FUNCTION_NAME); \
	\
	if(!original_fn_ptr) \
		printf("unable to get original ptr\n"); \
	\
	no_intercept = initial_no_intercept; \
	\
	if (return_error) \
	{ \
		errno = return_errno; \
		__asm__ ("pop %r9"); \
		__asm__ ("pop %r8"); \
		__asm__ ("pop %rdx"); \
		__asm__ ("pop %rcx"); \
		__asm__ ("pop %rbx"); \
		__asm__ ("pop %rsi"); \
		__asm__ ("pop %rdi"); \
		__asm__ ("pop %r10"); \
		__asm__ ("pop %r11"); \
		__asm__ ("pop %r12"); \
		__asm__ ("pop %r13"); \
		__asm__ ("nop" : : "a"(return_code)); \
		return; \
	} \
	else if (call_original) \
	{ \
		__asm__ ("pop %r9"); \
		__asm__ ("pop %r8"); \
		__asm__ ("pop %rdx"); \
		__asm__ ("pop %rcx"); \
		__asm__ ("pop %rbx"); \
		__asm__ ("pop %rsi"); \
		__asm__ ("pop %rdi"); \
		__asm__ ("pop %r10"); \
		__asm__ ("pop %r11"); \
		__asm__ ("pop %r12"); \
		__asm__ ("pop %r13"); \
		__asm__ ("leave"); \
		__asm__ ("jmp *%%rax" : : "a"(original_fn_ptr)); \
	} \
}


#define STUB_VAR_DECL \
int only_trace, trace_fd; /* legacy, unused */ \
int log_fd, replay_fd; \
int init_done; \
void* exe_start, *exe_end; \
time_t tt;
int injx_count;

